import os
import openai
import streamlit as st
import json
import fun_use
import funcs
openai.api_base = "https://yourapibase/v1"
openai.api_key = "yourapikey"
# 改成你自己的openai的api_base和api_key
st.header("MS08067小助手")
user_input = st.text_area('提出一个问题:', '', height=150)
st.markdown("""
    <style>
    .dashed-box {
        border: 2px dashed #FFA500;
        border-radius: 10px;
        padding: 10px;
        margin: 10px 0;
    }
    </style>
    """, unsafe_allow_html=True)
if st.button('提交'):
    # 改成你自己的模型
    # model="gpt-3.5-turbo"
    params = dict(model="gpt-3.5-turbo", messages=[{"role": "user", "content": user_input}], functions=funcs.tools)

    response = openai.ChatCompletion.create(**params)
    if response.choices[0].message.function_call is None:
        querycontent = response.choices[0].message.content
        st.markdown(f'<div class="dashed-box">{querycontent}</div>', unsafe_allow_html=True)
    elif response.choices[0].message.function_call.name == "run_code":
        #print(response.choices[0].message.function_call.arguments)
        function_args = json.loads(response.choices[0].message.function_call.arguments)
        filecontent = str(function_args["filecontent"])
        fun_use.run_code(filecontent=filecontent)
        if os.path.exists("testfile"):
            file_path = "testfile"
            f = open(file_path, "rb")
            download_content = f.read()
            # 创建一个下载链接
            href = f'<div class="dashed-box">文件创建成功，请点击下面的链接下载！</div>'

            st.markdown(href, unsafe_allow_html=True)
            st.download_button('下载文件', download_content,file_name="testfile")
        else:
            st.error(f"文件创建失败")
    elif response.choices[0].message.function_call.name == "nmap_scan":
        #print(response.choices[0].message.function_call.arguments)
        function_args = json.loads(response.choices[0].message.function_call.arguments)
        target = str(function_args["target"])
        #ports = str(function_args["ports"])
        querycontent = fun_use.nmap_scan(target=target)
        st.markdown(f'<div class="dashed-box">{querycontent}</div>', unsafe_allow_html=True)
